{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import covalent as ct"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Hello world"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Construct tasks as \"electrons\"\n",
    "@ct.electron\n",
    "def join_words(a, b):\n",
    "    return \", \".join([a, b])\n",
    "\n",
    "@ct.electron\n",
    "def excitement(a):\n",
    "    return f\"{a}!\"\n",
    "\n",
    "# Construct a workflow of tasks\n",
    "@ct.lattice\n",
    "def simple_workflow(a, b):\n",
    "    phrase = join_words(a, b)\n",
    "    return excitement(phrase)\n",
    "\n",
    "# Dispatch the workflow\n",
    "dispatch_id = ct.dispatch(simple_workflow)(\"Hello\", \"World\")\n",
    "result = ct.get_result(dispatch_id=dispatch_id, wait=True)\n",
    "assert result.result == 'Hello, World!'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Construct an electron"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from math import sqrt\n",
    "\n",
    "@ct.electron\n",
    "def identity(x):\n",
    "    return x\n",
    "\n",
    "@ct.electron\n",
    "def quadrature(x, y):\n",
    "    return sqrt(x**2 + y**2)\n",
    "\n",
    "# Check if identity is created\n",
    "identity_exists = True if 'identity' in locals() else False\n",
    "assert identity_exists == True\n",
    "\n",
    "# Check if quadrature is created\n",
    "quadrature_exists = True if 'quadrature' in locals() else False\n",
    "assert quadrature_exists == True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Construct a lepton"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Node entrypoint(0) failed: \n",
      "Traceback (most recent call last):\n",
      "  File \"/home/venkat/covalent/covalent_dispatcher/_core/execution.py\", line 221, in _run_task\n",
      "    output, stdout, stderr = executor.execute(\n",
      "  File \"/home/venkat/covalent/covalent/executor/executor_plugins/local.py\", line 106, in execute\n",
      "    result = fn(*args, **kwargs)\n",
      "  File \"/home/venkat/covalent/covalent/_workflow/lepton.py\", line 93, in python_wrapper\n",
      "    module = importlib.import_module(self.library_name)\n",
      "  File \"/home/venkat/anaconda3/envs/covalent-qa/lib/python3.8/importlib/__init__.py\", line 127, in import_module\n",
      "    return _bootstrap._gcd_import(name[level:], package, level)\n",
      "  File \"<frozen importlib._bootstrap>\", line 1014, in _gcd_import\n",
      "  File \"<frozen importlib._bootstrap>\", line 991, in _find_and_load\n",
      "  File \"<frozen importlib._bootstrap>\", line 973, in _find_and_load_unlocked\n",
      "ModuleNotFoundError: No module named '/home/venkat/covalent/tests/qa/test_module'\n",
      "\n"
     ]
    },
    {
     "ename": "AssertionError",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAssertionError\u001b[0m                            Traceback (most recent call last)",
      "\u001b[1;32m/home/venkat/covalent/tests/qa/qa_requirements.ipynb Cell 7'\u001b[0m in \u001b[0;36m<cell line: 23>\u001b[0;34m()\u001b[0m\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2Bdevbox/home/venkat/covalent/tests/qa/qa_requirements.ipynb#ch0000014vscode-remote?line=20'>21</a>\u001b[0m result \u001b[39m=\u001b[39m ct\u001b[39m.\u001b[39mdispatch_sync(workflow)(\u001b[39m2\u001b[39m)\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2Bdevbox/home/venkat/covalent/tests/qa/qa_requirements.ipynb#ch0000014vscode-remote?line=21'>22</a>\u001b[0m \u001b[39mprint\u001b[39m(result\u001b[39m.\u001b[39merror)\n\u001b[0;32m---> <a href='vscode-notebook-cell://ssh-remote%2Bdevbox/home/venkat/covalent/tests/qa/qa_requirements.ipynb#ch0000014vscode-remote?line=22'>23</a>\u001b[0m \u001b[39massert\u001b[39;00m result\u001b[39m.\u001b[39mstatus\u001b[39m.\u001b[39mSTATUS \u001b[39m==\u001b[39m \u001b[39m'\u001b[39m\u001b[39mCOMPLETED\u001b[39m\u001b[39m'\u001b[39m\n",
      "\u001b[0;31mAssertionError\u001b[0m: "
     ]
    }
   ],
   "source": [
    "import os\n",
    "\n",
    "test_module = \"\"\"\n",
    "def entrypoint(x: int, y: int) -> int:\n",
    "    return x + y\n",
    "\"\"\"\n",
    "with open(\"test_module.py\", \"w\") as f:\n",
    "    f.write(test_module)\n",
    "\n",
    "library_path = os.path.join(os.getcwd(), \"test_module\")\n",
    "task = ct.Lepton(\n",
    "    language = \"python\",\n",
    "    library_name = library_path,\n",
    "    function_name = \"entrypoint\",\n",
    ")\n",
    "\n",
    "@ct.lattice\n",
    "def workflow(x: int) -> int:\n",
    "    return task(x, 1)\n",
    "\n",
    "result = ct.dispatch_sync(workflow)(2)\n",
    "print(result.error)\n",
    "assert result.status.STATUS == 'COMPLETED'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Create a Lattice"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "@ct.lattice\n",
    "def workflow(x):\n",
    "    return x**2\n",
    "\n",
    "lattice_exists = True if 'workflow' in locals() else False\n",
    "assert lattice_exists == True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Visualize a lattice"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "@ct.electron\n",
    "def identity(x):\n",
    "    return x\n",
    "\n",
    "@ct.electron\n",
    "def square(x):\n",
    "    return x * x\n",
    "\n",
    "@ct.lattice\n",
    "def workflow(a):\n",
    "    val_1 = identity(x=a)\n",
    "    return square(x=val_1)\n",
    "\n",
    "dispatch_id = ct.dispatch(workflow)(a = 1)\n",
    "\n",
    "#from IPython import display\n",
    "#display.Image(\"images/visualize_lattice.png\", height=\"20pc\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Add constraints to lattice"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "@ct.electron\n",
    "def identity(x):\n",
    "    return x\n",
    "\n",
    "@ct.lattice(executor='local')\n",
    "def workflow(a):\n",
    "    return identity(x = a)\n",
    "\n",
    "assert workflow.metadata['executor'] == 'local'c"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Add Leptons in C"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "00000000000010f9 T test_entry\n",
      "(3, 5)\n"
     ]
    }
   ],
   "source": [
    "from ctypes import POINTER, c_int32\n",
    "import os\n",
    "\n",
    "c_source = \"\"\"\n",
    "#include \"test.h\"\n",
    "\n",
    "void test_entry(int x, int *y, int *z)\n",
    "{\n",
    "        *y += x;\n",
    "        *z = 5;\n",
    "}\n",
    "\"\"\"\n",
    "\n",
    "c_header = \"\"\"\n",
    "void test_entry(int x, int *y, int *z);\n",
    "\"\"\"\n",
    "\n",
    "with open(\"test.c\", \"w\") as f:\n",
    "    f.write(c_source)\n",
    "\n",
    "with open(\"test.h\", \"w\") as f:\n",
    "    f.write(c_header)\n",
    "\n",
    "# Compile C code\n",
    "!gcc -shared -fPIC -o libtest.so test.c\n",
    "\n",
    "!nm -D libtest.so | grep 'T test_entry'\n",
    "\n",
    "library_path = os.path.join(os.getcwd(),\"libtest.so\")\n",
    "task = ct.Lepton(\n",
    "    language = \"C\",\n",
    "    library_name = library_path,\n",
    "    function_name = \"test_entry\",\n",
    "    argtypes = [\n",
    "        (c_int32, ct.Lepton.INPUT),\n",
    "        (POINTER(c_int32), ct.Lepton.INPUT_OUTPUT),\n",
    "        (POINTER(c_int32), ct.Lepton.OUTPUT)\n",
    "    ]\n",
    ")\n",
    "\n",
    "@ct.lattice\n",
    "def workflow(x: int, y: int) -> int:\n",
    "    return task(x, y)\n",
    "\n",
    "result = ct.dispatch_sync(workflow)(1, 2)\n",
    "\n",
    "print(result.result)\n",
    "assert result.result == (3, 5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Construct a lepton in Bash"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "Language 'bash' is not supported.",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[1;32m/home/venkat/covalent/tests/qa/qa_requirements.ipynb Cell 17'\u001b[0m in \u001b[0;36m<cell line: 49>\u001b[0;34m()\u001b[0m\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2Bdevbox/home/venkat/covalent/tests/qa/qa_requirements.ipynb#ch0000030vscode-remote?line=45'>46</a>\u001b[0m \u001b[39mwith\u001b[39;00m \u001b[39mopen\u001b[39m(\u001b[39m\"\u001b[39m\u001b[39m/tmp/covalent_notify.sh\u001b[39m\u001b[39m\"\u001b[39m, \u001b[39m\"\u001b[39m\u001b[39mw\u001b[39m\u001b[39m\"\u001b[39m) \u001b[39mas\u001b[39;00m f:\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2Bdevbox/home/venkat/covalent/tests/qa/qa_requirements.ipynb#ch0000030vscode-remote?line=46'>47</a>\u001b[0m     f\u001b[39m.\u001b[39mwrite(bash_source)\n\u001b[0;32m---> <a href='vscode-notebook-cell://ssh-remote%2Bdevbox/home/venkat/covalent/tests/qa/qa_requirements.ipynb#ch0000030vscode-remote?line=48'>49</a>\u001b[0m task \u001b[39m=\u001b[39m ct\u001b[39m.\u001b[39;49mLepton(\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2Bdevbox/home/venkat/covalent/tests/qa/qa_requirements.ipynb#ch0000030vscode-remote?line=49'>50</a>\u001b[0m     \u001b[39m\"\u001b[39;49m\u001b[39mbash\u001b[39;49m\u001b[39m\"\u001b[39;49m,\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2Bdevbox/home/venkat/covalent/tests/qa/qa_requirements.ipynb#ch0000030vscode-remote?line=50'>51</a>\u001b[0m     \u001b[39m\"\u001b[39;49m\u001b[39m/tmp/covalent_notify.sh\u001b[39;49m\u001b[39m\"\u001b[39;49m,\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2Bdevbox/home/venkat/covalent/tests/qa/qa_requirements.ipynb#ch0000030vscode-remote?line=51'>52</a>\u001b[0m     \u001b[39m\"\u001b[39;49m\u001b[39mnotify_slack\u001b[39;49m\u001b[39m\"\u001b[39;49m\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2Bdevbox/home/venkat/covalent/tests/qa/qa_requirements.ipynb#ch0000030vscode-remote?line=52'>53</a>\u001b[0m )\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2Bdevbox/home/venkat/covalent/tests/qa/qa_requirements.ipynb#ch0000030vscode-remote?line=54'>55</a>\u001b[0m \u001b[39m@ct\u001b[39m\u001b[39m.\u001b[39mlattice\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2Bdevbox/home/venkat/covalent/tests/qa/qa_requirements.ipynb#ch0000030vscode-remote?line=55'>56</a>\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mworkflow\u001b[39m(display_name: \u001b[39mstr\u001b[39m, message: \u001b[39mstr\u001b[39m) \u001b[39m-\u001b[39m\u001b[39m>\u001b[39m \u001b[39mstr\u001b[39m:\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2Bdevbox/home/venkat/covalent/tests/qa/qa_requirements.ipynb#ch0000030vscode-remote?line=56'>57</a>\u001b[0m     \u001b[39mreturn\u001b[39;00m task(display_name, message)\n",
      "File \u001b[0;32m~/covalent/covalent/_workflow/lepton.py:79\u001b[0m, in \u001b[0;36mLepton.__init__\u001b[0;34m(self, language, library_name, function_name, argtypes, executor)\u001b[0m\n\u001b[1;32m     <a href='file:///home/venkat/covalent/covalent/_workflow/lepton.py?line=75'>76</a>\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39margtypes \u001b[39m=\u001b[39m [(arg[\u001b[39m0\u001b[39m]\u001b[39m.\u001b[39m\u001b[39m__name__\u001b[39m, arg[\u001b[39m1\u001b[39m]) \u001b[39mfor\u001b[39;00m arg \u001b[39min\u001b[39;00m argtypes]\n\u001b[1;32m     <a href='file:///home/venkat/covalent/covalent/_workflow/lepton.py?line=77'>78</a>\u001b[0m \u001b[39m# Assign the wrapper below as the task's callable function\u001b[39;00m\n\u001b[0;32m---> <a href='file:///home/venkat/covalent/covalent/_workflow/lepton.py?line=78'>79</a>\u001b[0m \u001b[39msuper\u001b[39m()\u001b[39m.\u001b[39m\u001b[39m__init__\u001b[39m(\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mwrap_task())\n\u001b[1;32m     <a href='file:///home/venkat/covalent/covalent/_workflow/lepton.py?line=80'>81</a>\u001b[0m \u001b[39m# Assign metadata defaults\u001b[39;00m\n\u001b[1;32m     <a href='file:///home/venkat/covalent/covalent/_workflow/lepton.py?line=81'>82</a>\u001b[0m \u001b[39msuper\u001b[39m()\u001b[39m.\u001b[39mset_metadata(\u001b[39m\"\u001b[39m\u001b[39mexecutor\u001b[39m\u001b[39m\"\u001b[39m, executor)\n",
      "File \u001b[0;32m~/covalent/covalent/_workflow/lepton.py:201\u001b[0m, in \u001b[0;36mLepton.wrap_task\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    <a href='file:///home/venkat/covalent/covalent/_workflow/lepton.py?line=198'>199</a>\u001b[0m     wrapper \u001b[39m=\u001b[39m c_wrapper\n\u001b[1;32m    <a href='file:///home/venkat/covalent/covalent/_workflow/lepton.py?line=199'>200</a>\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[0;32m--> <a href='file:///home/venkat/covalent/covalent/_workflow/lepton.py?line=200'>201</a>\u001b[0m     \u001b[39mraise\u001b[39;00m \u001b[39mValueError\u001b[39;00m(\u001b[39mf\u001b[39m\u001b[39m\"\u001b[39m\u001b[39mLanguage \u001b[39m\u001b[39m'\u001b[39m\u001b[39m{\u001b[39;00m\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mlanguage\u001b[39m}\u001b[39;00m\u001b[39m'\u001b[39m\u001b[39m is not supported.\u001b[39m\u001b[39m\"\u001b[39m)\n\u001b[1;32m    <a href='file:///home/venkat/covalent/covalent/_workflow/lepton.py?line=202'>203</a>\u001b[0m \u001b[39m# Attribute translation\u001b[39;00m\n\u001b[1;32m    <a href='file:///home/venkat/covalent/covalent/_workflow/lepton.py?line=203'>204</a>\u001b[0m wrapper\u001b[39m.\u001b[39m\u001b[39m__name__\u001b[39m \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mfunction_name\n",
      "\u001b[0;31mValueError\u001b[0m: Language 'bash' is not supported."
     ]
    }
   ],
   "source": [
    "bash_source = \"\"\"#! /bin/bash\n",
    "\n",
    "set -eu -o pipefail\n",
    "\n",
    "# Below is a bot token for a Slack bot with the following OAuth scopes:\n",
    "# - users:read\n",
    "# - chat:write\n",
    "# - groups:write\n",
    "# - im:write\n",
    "# - mpim:write\n",
    "\n",
    "# You'll need to replace this with your own bot token\n",
    "token=\"xoxb-abcdef-0123456789-abcdef0123456789\"\n",
    "\n",
    "notify_slack(){\n",
    "  if [[ \"$#\" -lt 2 ]] ; then\n",
    "    echo \"Arguments are [display_name] [message].\"\n",
    "    return 1\n",
    "  fi\n",
    "\n",
    "  display_name=\"$1\"\n",
    "  message=\"$2\"\n",
    "\n",
    "  # Retrieve the user ID from a given display name\n",
    "  id=`curl -s -X POST \"https://slack.com/api/users.list\" \\\\\n",
    "    -H \"accept: application/json\" \\\\\n",
    "    -d token=\"$token\" | jq -r '.members[] | select(.profile.display_name==\"'$display_name'\").id'`\n",
    "  echo $id\n",
    "\n",
    "  # Retrieve a channel ID for a direct message to the user\n",
    "  channel=`curl -s -X POST \"https://slack.com/api/conversations.open\" \\\\\n",
    "    -H \"accept: application/json\" \\\\\n",
    "    -d token=\"$token\" \\\\\n",
    "    -d users=\"$id\" | jq -r '.channel.id'`\n",
    "\n",
    "  # Post the message to the user in Slack\n",
    "  curl -s -X POST \"https://slack.com/api/chat.postMessage\" \\\\\n",
    "    -H \"accept: application/json\" \\\\\n",
    "    -d token=\"$token\" \\\\\n",
    "    -d channel=\"$channel\" \\\\\n",
    "    -d text=\"$message\" \\\\\n",
    "    -d as_user=true\n",
    "}\n",
    "\"\"\"\n",
    "\n",
    "with open(\"/tmp/covalent_notify.sh\", \"w\") as f:\n",
    "    f.write(bash_source)\n",
    "\n",
    "task = ct.Lepton(\n",
    "    \"bash\",\n",
    "    \"/tmp/covalent_notify.sh\",\n",
    "    \"notify_slack\"\n",
    ")\n",
    "\n",
    "@ct.lattice\n",
    "def workflow(display_name: str, message: str) -> str:\n",
    "    return task(display_name, message)\n",
    "\n",
    "response = ct.dispatch_sync(workflow)(\n",
    "    display_name=\"will\",\n",
    "    message=\"Your workflow has successfully completed!\"\n",
    ")\n",
    "\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "96665897554d405553c080f7a3b66ad582de30e62526c6c9d18f08acffe92e34"
  },
  "kernelspec": {
   "display_name": "Python 3.8.13 ('covalent-qa')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
